"""
基于 Pinecone 的在线 KNN 推理 WebApp：
- 复用 `optimal_knn_webapp.py` 的预处理，将任意手写图像转为 digits(8x8) 风格 64 维向量；
- 每次预测均调用 Pinecone 索引进行向量搜索（top_k=11），再多数投票给出结果；
- 首次调用可能较慢（索引冷启动/网络），后续会快一些；
- 需要在环境变量中提供 PINECONE_API_KEY；索引名默认为 "mnist-index"。
"""

from __future__ import annotations

import base64
import io
import os
from collections import Counter
from typing import Optional

import gradio as gr
import numpy as np
from PIL import Image

try:
    from dotenv import load_dotenv  # type: ignore

    load_dotenv()  # noqa: F401
except Exception:
    pass

from pinecone import Pinecone  # type: ignore


INDEX_NAME = os.getenv("PINECONE_INDEX", "mnist-index")
API_KEY = os.getenv("PINECONE_API_KEY")
TOP_K = int(os.getenv("K", "11"))


def preprocess_to_digits_like(img: Image.Image) -> np.ndarray:
    img = img.convert("L")
    np_img = 255.0 - np.array(img, dtype=np.float32)
    img_small = Image.fromarray(np.clip(np_img, 0, 255).astype(np.uint8)).resize((8, 8), Image.BILINEAR)
    arr = np.array(img_small, dtype=np.float32)
    arr = arr / 255.0 * 16.0
    return arr.reshape(1, -1)


def to_pil(image) -> Image.Image:
    # 兼容 gradio Sketchpad/Image 的多种返回类型
    if image is None:
        raise gr.Error("请先在画布上写一个数字。")

    if isinstance(image, dict):
        candidate = None
        for key in ("image", "composite", "background"):
            if key in image and image[key] is not None:
                candidate = image[key]
                break
        if candidate is None and "layers" in image and image["layers"]:
            layers = image["layers"]
            norm_layers = []
            for lyr in layers:
                if isinstance(lyr, list):
                    lyr = np.array(lyr)
                norm_layers.append(np.asarray(lyr))
            h = max(l.shape[0] for l in norm_layers)
            w = max(l.shape[1] for l in norm_layers)
            canvas = Image.new("RGBA", (w, h), (255, 255, 255, 255))
            for arr in norm_layers:
                if arr.dtype != np.uint8:
                    maxv = float(np.max(arr)) if arr.size else 1.0
                    if maxv <= 1.0:
                        arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
                    else:
                        arr = np.clip(arr, 0, 255).astype(np.uint8)
                if arr.ndim == 2:
                    pil = Image.fromarray(arr, mode="L").convert("RGBA")
                elif arr.ndim == 3:
                    if arr.shape[-1] == 4:
                        pil = Image.fromarray(arr, mode="RGBA")
                    else:
                        pil = Image.fromarray(arr[:, :, :3], mode="RGB").convert("RGBA")
                else:
                    continue
                canvas.alpha_composite(pil)
            return canvas.convert("RGB")
        else:
            image = candidate

    if isinstance(image, list):
        image = np.array(image)

    if isinstance(image, (bytes, bytearray)):
        return Image.open(io.BytesIO(image)).convert("RGB")
    if isinstance(image, str) and image.startswith("data:image"):
        b64 = image.split(",", 1)[1]
        return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
    if isinstance(image, np.ndarray):
        arr = image
        if arr.dtype != np.uint8:
            maxv = float(np.max(arr)) if arr.size else 1.0
            if maxv <= 1.0:
                arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
            else:
                arr = np.clip(arr, 0, 255).astype(np.uint8)
        if arr.ndim == 2:
            return Image.fromarray(arr, mode="L")
        if arr.ndim == 3:
            if arr.shape[-1] == 4:
                pil = Image.fromarray(arr, mode="RGBA")
                bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
                pil = Image.alpha_composite(bg, pil).convert("RGB")
                return pil
            return Image.fromarray(arr[:, :, :3], mode="RGB")
        raise gr.Error("无法识别的图像维度。")
    if isinstance(image, Image.Image):
        return image
    raise gr.Error("无法识别的图像格式。")


class PineconeClientHolder:
    _pc = None
    _index = None

    @classmethod
    def get_index(cls):
        if cls._index is None:
            if not API_KEY:
                raise gr.Error("未配置 PINECONE_API_KEY 环境变量。")
            cls._pc = Pinecone(api_key=API_KEY)
            cls._index = cls._pc.Index(INDEX_NAME)
        return cls._index


def predict_fn(image) -> int:
    pil = to_pil(image)
    vec = preprocess_to_digits_like(pil)[0].tolist()

    index = PineconeClientHolder.get_index()
    res = index.query(vector=vec, top_k=TOP_K, include_metadata=True)
    matches = res.get("matches") if isinstance(res, dict) else getattr(res, "matches", [])

    labels = []
    for m in matches:
        md = m.get("metadata", {}) if isinstance(m, dict) else getattr(m, "metadata", {})
        if md and "label" in md:
            labels.append(int(md["label"]))
    if not labels:
        raise gr.Error("未从 Pinecone 返回有效邻居，请确认索引已载入数据。")

    pred = Counter(labels).most_common(1)[0][0]
    return int(pred)


with gr.Blocks(title="KNN 手写数字识别 (Pinecone)") as demo:
    gr.Markdown("# 使用 Pinecone 的 KNN 手写数字识别\n在画布上写 0~9 的数字，点击提交后将在 Pinecone 上进行近邻搜索。")

    with gr.Row():
        try:
            sketch = gr.Sketchpad(label="手写板", image_mode="L", brush=20, height=256, width=256)
            input_comp = sketch
        except Exception:
            img = gr.Image(label="手写板(兼容模式)", type="pil", tool=None, image_mode="L", height=256, width=256)
            input_comp = img

        out = gr.Number(label="预测结果", precision=0)

    btn = gr.Button("提交")
    btn.click(predict_fn, inputs=input_comp, outputs=out)

if __name__ == "__main__":
    import os as _os

    share = _os.getenv("GRADIO_SHARE", "0").lower() in {"1", "true", "yes", "y"}
    port_env = _os.getenv("PORT") or _os.getenv("GRADIO_PORT") or _os.getenv("SERVER_PORT")
    try:
        port = int(port_env) if port_env else None
    except Exception:
        port = None
    server_name = _os.getenv("SERVER_NAME", "0.0.0.0")

    demo.launch(share=share, server_name=server_name, server_port=port)
